#TODO 导入相关的Python库或Python模块
import torch
import torch.nn as nn  
import torch.nn.functional as F  
import torchvision.transforms as transforms  
from torch.utils.data import DataLoader, Dataset  
from PIL import Image  
import os  
import matplotlib.pyplot as plt
import os
from PIL import Image
from torch.utils.data import Dataset
  
# 模型定义
class CNNNet(nn.Module):
    # 模型初始化
    def __init__(self):
        super(CNNNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(500, 50)  # 注意这里可能需要调整以适应你的图像尺寸
        self.fc2 = nn.Linear(50, 4)

    # 模型的前向传播
    def forward(self, x):  
        x = F.relu(F.max_pool2d(self.conv1(x), 2))  
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))  
        x = x.view(-1, 500)  # 注意这里可能需要调整  
        x = F.relu(self.fc1(x))  
        x = F.dropout(x, training=self.training)  
        x = self.fc2(x)  
        return F.log_softmax(x, dim=1)  
  

device = torch.device("cpu")#指定模型将在CPU上运行
model = CNNNet()
#TODO 本题文件的加载路径不唯一，需根据实际情况进行修改
model.load_state_dict(torch.load('D:\\横向1\\hengxiang\\2\\pre.pt'))  
model.eval()  
  
# 数据预处理  
transform = transforms.Compose([
    transforms.Resize((32, 32)),            # 将图像大小调整为32像素宽和32像素高
    transforms.ToTensor(),                   # 将PIL图像或numpy数组转换为torch.FloatTensor
    transforms.Normalize((0.5,), (0.5,))    # 对灰度图进行标准化，即将数据转换到一个特定的均值和标准差范围
])  




#TODO 创建自定义数据集
class TestDataset(Dataset):
    # TestDataset类的构造函数
    def __init__(self, root_dir, transform=None):  
        self.root_dir = root_dir  
        self.transform = transform  
        self.images = []  #初始化两个空列表，用于存储图像的路径和对应的标签
        self.labels = []  
        self.label_to_int = {'A': 0, 'B': 1, 'C': 2, 'D': 3}  # 定义标签到整数的映射  
  
        #TODO 加载图像和标签
        for label in ['A', 'B', 'C', 'D']:  
            class_dir = os.path.join(root_dir, label)
            # TODO（1） 遍历目录中所有图像，将每个图像的路径添加到self.images列表中，并将对应的标签添加到self.labels列表中
            for filename in os.listdir(class_dir):  
                img_path = os.path.join(class_dir, filename)  
                self.images.append(img_path)  
                self.labels.append(label)  # 这里仍然是字符串，但稍后会转换  

    # TODO 获取数据集大小
    def __len__(self):
        # TODO（1） 返回数据集中图像的数量
        return len(self.images)  

    # 从数据集中获取索引加载和预处理数据集中的单个样本
    def __getitem__(self, idx):  
        img_path = self.images[idx]  
        image = Image.open(img_path).convert('L')  # 转换为灰度图  

        # 应用图像变换
        if self.transform:  
            image = self.transform(image)
        #TODO（1） 使用self.label_to_int字典将标签字符串转换为整数
        label = self.label_to_int[self.labels[idx]]
        # TODO（2）# 创建一维张量
        label_tensor = torch.tensor([label], dtype=torch.long)
  
        return image, label_tensor



#TODO 加载测试集，本题文件的加载路径不唯一，需根据实际情况进行修改
test_dataset = TestDataset('D:\\横向1\\hengxiang\\dataset\\test', transform=transform)  
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)  


for images, labels in test_loader:  
    print(type(images), type(labels))  
    # 其他代码...

# 预测并统计结果  
correct_counts = {label: 0 for label in range(4)}  # 使用整数键 0, 1, 2, 3  
total_counts = {label: 0 for label in range(4)}    # 同上  
  
with torch.no_grad():  
    for images, labels in test_loader:  
        images, labels = images.to(device), labels.to(device)

        outputs = model(images)     # 将图像数据传递给模型进行前向传播
        _, predicted = torch.max(outputs, 1)  
          
        # 直接遍历 labels 和 predicted
        for true_label, pred_label in zip(labels.cpu(), predicted.cpu()):  # 转到 CPU 以使用 .item()
            # TODO（1） 对比每个真实标签和预测标签是否相等。如果相等，则表示预测正确，并更新对应类别的正确计数器
            if true_label.item() == pred_label.item():  
                correct_counts[true_label.item()] += 1
            # TODO（2） 更新对应类别的总预测数
            total_counts[true_label.item()] += 1


#TODO 绘制正确率柱状图
labels = ['A', 'B', 'C', 'D']  
correct_rates = [correct_counts.get(label, 0) / max(total_counts.get(label, 0), 1) * 100 for label in range(4)]  

#TODO（1） 创建一个新的图形窗口，宽和高参数分别为8，6
plt.figure(figsize=(8, 6))
#TODO（2） 使用plt的bar函数绘制条形图，条形的颜色为skyblue
plt.bar(labels, correct_rates, color='skyblue')
#TODO（3） 设置图表的标题为’Accuracy Rates (%)‘，x轴和y轴的标签分别为‘Labels’，‘Accuracy’
plt.title('Accuracy Rates (%)')  
plt.xlabel('Labels')  
plt.ylabel('Accuracy')  
plt.tight_layout()  
plt.show()
